#pragma once

#include <cmath>
#include <cstdlib>
#include <memory>
#include <mutex>
#include <sstream>
#include <vector>
#include <shared_mutex>

#include "KVDatabase.h"

template <typename K, typename V>
class SkipList : public KVDatabase<K, V> {
private:
    // 节点结构
    struct Node {
        K key;
        V value;
        std::unique_ptr<Node[]> forward;

        Node(K k, V v, int level) : key(k), value(v), forward(new Node[level + 1]) {}
    };

    int max_level_;
    int skip_list_level_;
    int element_count_;
    std::shared_mutex mtx_;
    std::unique_ptr<Node> header_;

public:
    explicit SkipList(int max_level = 16)
        : max_level_(max_level),
          skip_list_level_(0),
          element_count_(0),
          header_(std::make_unique<Node>(K{}, V{}, max_level)) {}

    ~SkipList() override {
        clear();
    }

    // 插入或更新键值对
    bool put(const K& key, const V& value) override {
        std::unique_lock lock(mtx_);

        Node* update[max_level_ + 1] = {nullptr};
        Node* current = header_.get();

        for (int i = skip_list_level_; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i].get();
            }
            update[i] = current;
        }

        current = current->forward[0].get();

        if (current && current->key == key) {
            current->value = value; // 更新
            return true;
        }

        int random_level = get_random_level();
        if (random_level > skip_list_level_) {
            for (int i = skip_list_level_ + 1; i <= random_level; ++i) {
                update[i] = header_.get();
            }
            skip_list_level_ = random_level;
        }

        Node* new_node = new Node(key, value, random_level);

        for (int i = 0; i <= random_level; ++i) {
            new_node->forward[i] = std::move(update[i]->forward[i]);
            update[i]->forward[i].reset(new_node);
        }

        ++element_count_;
        return true;
    }

    // 查询键值对
    bool get(const K& key, V& value) const override {
        std::shared_lock lock(mtx_);

        Node* current = header_.get();
        for (int i = skip_list_level_; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i].get();
            }
        }

        current = current->forward[0].get();
        if (current && current->key == key) {
            value = current->value;
            return true;
        }

        return false;
    }

    // 判断键是否存在
    bool exists(const K& key) const override {
        std::shared_lock lock(mtx_);

        Node* current = header_.get();
        for (int i = skip_list_level_; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i].get();
            }
        }

        current = current->forward[0].get();
        return current && current->key == key;
    }

    // 删除键值对
    bool remove(const K& key) override {
        std::unique_lock lock(mtx_);

        Node* update[max_level_ + 1] = {nullptr};
        Node* current = header_.get();

        for (int i = skip_list_level_; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i].get();
            }
            update[i] = current;
        }

        current = current->forward[0].get();
        if (!current || current->key != key) {
            return false;
        }

        for (int i = 0; i <= skip_list_level_; ++i) {
            if (update[i]->forward[i].get() != current) break;
            update[i]->forward[i].reset(current->forward[i].release());
        }

        delete current;

        while (skip_list_level_ > 0 && !header_->forward[skip_list_level_]) {
            --skip_list_level_;
        }

        --element_count_;
        return true;
    }

    // 序列化：将当前数据库状态写入 data_str
    bool serialize(std::string& data_str) const override {
        std::unique_lock lock(mtx_);

        std::ostringstream oss;
        boost::archive::text_oarchive oa(oss);

        std::vector<std::pair<K, V>> entries;
        Node* current = header_->forward[0].get();
        while (current) {
            entries.emplace_back(current->key, current->value);
            current = current->forward[0].get();
        }

        oa << entries;
        data_str = oss.str();
        return true;
    }

    // 反序列化：从 data_str 加载数据库状态
    bool deserialize(const std::string& data_str) override {
        std::unique_lock lock(mtx_);

        clear();

        std::istringstream iss(data_str);
        boost::archive::text_iarchive ia(iss);

        std::vector<std::pair<K, V>> entries;
        ia >> entries;

        for (const auto& [key, value] : entries) {
            put(key, value);
        }

        return true;
    }

    // 获取当前键值对数量
    size_t size() const override {
        std::shared_lock lock(mtx_);
        return element_count_;
    }

private:
    int get_random_level() const {
        int level = 0;
        while (rand() % 2 == 0 && level < max_level_) {
            ++level;
        }
        return level;
    }

    void clear() {
        Node* current = header_->forward[0].get();
        while (current) {
            Node* next = current->forward[0].get();
            delete current;
            current = next;
        }
        header_->forward[0].reset(nullptr);
        skip_list_level_ = 0;
        element_count_ = 0;
    }
};